
import numpy as np
import gradio as gr
import joblib
from skimage.transform import resize
import matplotlib.pyplot as plt

# 加载保存的KNN模型
model = joblib.load('best_knn_model.pkl')

# 推断模型需要的特征维度（用于决定缩放尺寸）
try:
    _n_features = int(getattr(model, 'n_features_in_', 64))
except Exception:
    _n_features = 64

_side = int(np.sqrt(_n_features)) if int(np.sqrt(_n_features)) ** 2 == _n_features else 8
_side = 8 if _side not in (8, 28) else _side  # 仅支持常见的 8x8 或 28x28，默认 8
_scale_max = 16 if _side == 8 else 255
print(f"[INFO] model expects {_n_features} features -> target size: {_side}x{_side}, scale_max={_scale_max}")

def predict_digit(image, invert=False):
    if image is None:
        return "请在画布上绘制一个数字。"

    # 将输入统一为 numpy 数组
    try:
        # 兼容某些版本返回 dict 的情况
        if isinstance(image, dict):
            if 'image' in image and image['image'] is not None:
                image = image['image']
            elif 'composite' in image and image['composite'] is not None:
                image = image['composite']
            else:
                return "请在画布上绘制一个数字。"
        if not isinstance(image, np.ndarray):
            image = np.array(image)
    except Exception as e:
        return f"无法解析输入图像: {e}"

    # 若是 0 维对象数组或空图，直接提示重画
    if not isinstance(image, np.ndarray) or image.size == 0 or getattr(image, 'ndim', 0) < 2:
        return "请在画布上绘制一个更清晰的数字。"

    # 打印原始图像的形状和类型
    try:
        print(f"Original image shape: {image.shape}, dtype: {image.dtype}")
    except Exception:
        pass

    # 确保图像是 2D 的灰度图像
    if image.ndim == 3:
        # RGBA 或 RGB
        if image.shape[2] == 4:
            image_gray = image[:, :, :3].mean(axis=2)
        else:
            image_gray = image.mean(axis=2)
    else:
        image_gray = image

    # 反色（可选）
    if invert:
        # 若值范围是 0..255，则先转到 0..1 再反色
        if image_gray.max() > 1.0:
            image_gray = image_gray / 255.0
        image_gray = 1.0 - image_gray

    # 归一化到 0..1（skimage.resize 处理浮点 0..1 更稳定）
    try:
        if image_gray.max() > 1.0:
            image_gray = image_gray / 255.0
    except Exception:
        return "请在画布上绘制一个数字（当前图像无效）。"

    # 自动反色：若整体较亮（白底黑字场景），则将笔迹置为高值
    try:
        if not invert and float(image_gray.mean()) > 0.5:
            image_gray = 1.0 - image_gray
            print('[INFO] auto invert applied')
    except Exception:
        pass

    # 打印灰度图像的形状和类型
    try:
        print(f"Grayscale image shape: {image_gray.shape}, dtype: {image_gray.dtype}, min={image_gray.min():.4f}, max={image_gray.max():.4f}")
    except Exception:
        pass

    # 提升准确率：裁剪前景、居中到正方形、再缩放到目标边长
    try:
        # 将笔迹视为高值（经过自动/手动反色后）
        thr = max(0.2, float(image_gray.mean()) * 0.7)
        mask = image_gray > thr
        if mask.sum() < 10:
            return "笔迹太淡或太小，请加粗并画大一些。"
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)
        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]
        cropped = image_gray[rmin:rmax+1, cmin:cmax+1]

        # 居中到正方形（填充背景 0）
        h, w = cropped.shape
        L = max(h, w)
        pad_vert = L - h
        pad_horz = L - w
        top = pad_vert // 2
        bottom = pad_vert - top
        left = pad_horz // 2
        right = pad_horz - left
        squared = np.pad(cropped, ((top, bottom), (left, right)), mode='constant', constant_values=0.0)

        image_resized = resize(squared, (_side, _side), anti_aliasing=True)
    except Exception as e:
        # 回退：直接缩放
        image_resized = resize(image_gray, (_side, _side), anti_aliasing=True)

    # 打印调整大小后的图像形状和类型
    try:
        print(f"Resized image shape: {image_resized.shape}, dtype: {image_resized.dtype}")
    except Exception:
        pass

    # 将像素值范围调整到 0-16（8x8）或 0-255（28x28）
    image_resized = np.clip(image_resized * _scale_max, 0, _scale_max)
    image_resized = image_resized.astype(int)

    # 打印调整范围后的图像
    try:
        print("Resized and scaled image:")
        print(image_resized)
    except Exception:
        pass

    # 展平图像
    image_flattened = image_resized.flatten()

    # 打印展平后的图像形状
    try:
        print(f"Flattened image shape: {image_flattened.shape}")
    except Exception:
        pass

    # 进行预测（异常兜底）
    try:
        prediction = model.predict([image_flattened])[0]
    except Exception as e:
        return f"预测失败: {e}"

    # 保存处理后的图像以供检查
    try:
        plt.imsave('processed_image.png', image_resized, cmap='gray')
    except Exception:
        pass

    return f"预测的数字是: {prediction}"

"""
优先提供“手绘画布”而非上传图片：
- 新版：Image(image_mode="L", source="canvas")
- 老版优先：gr.inputs.Sketchpad 或 gr.Sketchpad
- 最后兜底：Image()
"""

# 创建输入组件（优先支持手绘）
image_input = None

# 1) 老版本 inputs.Sketchpad（非常常见）
inputs_mod = getattr(gr, "inputs", None)
if inputs_mod is not None and hasattr(inputs_mod, "Sketchpad"):
    image_input = inputs_mod.Sketchpad()

# 2) 新版/中版：gr.Sketchpad
if image_input is None:
    Sketchpad = getattr(gr, "Sketchpad", None)
    if Sketchpad is not None:
        image_input = Sketchpad()

# 3) 新版 Image 画布
if image_input is None:
    try:
        image_input = gr.Image(image_mode="L", source="canvas")
    except TypeError:
        try:
            image_input = gr.Image(image_mode="L")
        except TypeError:
            image_input = gr.Image()

# 可选：反色开关（默认关闭）
invert_checkbox = gr.Checkbox(label="反色（将背景变黑，笔迹变白）", value=False)

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=[image_input, invert_checkbox],
    outputs="text",
    live=True,
    title="手写数字识别（手绘版）",
    description="在画布上手绘一个数字（可勾选反色适配），模型将实时预测结果。",
    allow_flagging="never"
)

# 启动Gradio接口
iface.launch()